#!/usr/bin/env python3
"""
RAG问答系统 Web API服务
提供RESTful API接口，支持单问题问答、批量问答和自定义参数问答

作者：AI Assistant
日期：2025-10-04
"""

from flask import Flask, request, jsonify, render_template, send_from_directory
from flask_cors import CORS
import os
import logging
from typing import List, Dict
from dotenv import load_dotenv
from rag_qa_system import RAGQASystem

# 加载环境变量
load_dotenv()

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 创建Flask应用
app = Flask(__name__)
CORS(app)  # 启用跨域支持

# 全局RAG系统实例
rag_system = None

# API配置 - 从环境变量读取
API_KEY = os.getenv("MODELARTS_API_KEY")
BASE_URL = os.getenv("MODELARTS_BASE_URL", "https://api.modelarts-maas.com/v1")
MODEL_NAME = os.getenv("MODELARTS_MODEL_NAME", "openpangu-pro-moe-72b")

# 验证必需的环境变量
if not API_KEY:
    logger.error("❌ 错误：未设置 MODELARTS_API_KEY 环境变量")
    logger.error("请创建 .env 文件并设置 MODELARTS_API_KEY，参考 .env.example 文件")
    raise ValueError("MODELARTS_API_KEY 环境变量未设置")


def initialize_rag_system():
    """初始化RAG系统"""
    global rag_system
    
    if rag_system is not None:
        return True
    
    try:
        logger.info("初始化RAG系统...")
        rag_system = RAGQASystem(
            api_key=API_KEY,
            base_url=BASE_URL,
            model_name=MODEL_NAME
        )
        
        if rag_system.load_vector_database():
            logger.info("✅ RAG系统初始化成功")
            return True
        else:
            logger.error("❌ 向量数据库加载失败")
            return False
            
    except Exception as e:
        logger.error(f"❌ RAG系统初始化失败: {str(e)}")
        return False


@app.route('/')
def index():
    """主页"""
    return render_template('index.html')


@app.route('/static/<path:path>')
def send_static(path):
    """静态文件服务"""
    return send_from_directory('static', path)


@app.route('/api/status', methods=['GET'])
def get_status():
    """获取系统状态"""
    try:
        if rag_system is None:
            return jsonify({
                'status': 'not_initialized',
                'message': '系统未初始化',
                'initialized': False
            })
        
        # 获取向量数据库统计信息
        collection = rag_system.vector_store._collection
        doc_count = collection.count()
        
        return jsonify({
            'status': 'ready',
            'message': '系统运行正常',
            'initialized': True,
            'model_name': MODEL_NAME,
            'document_count': doc_count,
            'vector_db_path': rag_system.vector_db_path
        })
        
    except Exception as e:
        logger.error(f"获取状态失败: {str(e)}")
        return jsonify({
            'status': 'error',
            'message': f'获取状态失败: {str(e)}',
            'initialized': False
        }), 500


@app.route('/api/ask', methods=['POST'])
def ask_question():
    """单问题问答接口"""
    try:
        # 检查系统是否初始化
        if rag_system is None:
            return jsonify({
                'success': False,
                'error': '系统未初始化，请先初始化系统'
            }), 503
        
        # 获取请求参数
        data = request.get_json()
        question = data.get('question', '').strip()
        
        if not question:
            return jsonify({
                'success': False,
                'error': '问题不能为空'
            }), 400
        
        # 获取可选参数
        max_docs = data.get('max_docs', 5)
        score_threshold = data.get('score_threshold', 1.0)
        
        # 调用RAG系统
        result = rag_system.ask_question(
            question=question,
            max_docs=max_docs,
            score_threshold=score_threshold
        )
        
        # 返回结果
        if result['answer']:
            return jsonify({
                'success': True,
                'question': result['question'],
                'answer': result['answer'],
                'retrieved_docs': result['retrieved_docs'],
                'sources': result['sources']
            })
        else:
            return jsonify({
                'success': False,
                'error': result.get('error', '未知错误')
            }), 500
            
    except Exception as e:
        logger.error(f"处理问题失败: {str(e)}")
        return jsonify({
            'success': False,
            'error': f'服务器错误: {str(e)}'
        }), 500


@app.route('/api/initialize', methods=['POST'])
def initialize_system():
    """初始化系统接口"""
    try:
        if initialize_rag_system():
            return jsonify({
                'success': True,
                'message': 'RAG系统初始化成功'
            })
        else:
            return jsonify({
                'success': False,
                'error': '系统初始化失败，请检查向量数据库和模型配置'
            }), 500
            
    except Exception as e:
        logger.error(f"初始化失败: {str(e)}")
        return jsonify({
            'success': False,
            'error': f'初始化失败: {str(e)}'
        }), 500


# 错误处理
@app.errorhandler(404)
def not_found(error):
    return jsonify({'error': '接口不存在'}), 404


@app.errorhandler(500)
def internal_error(error):
    return jsonify({'error': '服务器内部错误'}), 500


def main():
    """主函数"""
    print("🚀 RAG问答系统 Web服务")
    print("=" * 60)
    
    # 启动时初始化系统
    print("🔄 初始化RAG系统...")
    if initialize_rag_system():
        print("✅ 系统初始化成功")
    else:
        print("⚠️  系统初始化失败，但服务仍将启动")
        print("    您可以通过 /api/initialize 接口重新初始化")
    
    print()
    print("📡 服务已启动，访问地址:")
    print("   - 主页: http://localhost:5000")
    print("   - API文档: http://localhost:5000/api/status")
    print()
    print("💡 可用的API接口:")
    print("   - GET  /api/status        - 获取系统状态")
    print("   - POST /api/initialize    - 初始化系统")
    print("   - POST /api/ask           - 单问题问答")
    print()
    print("按 Ctrl+C 停止服务")
    print("=" * 60)
    
    # 启动Flask应用
    app.run(host='0.0.0.0', port=5000, debug=True)


if __name__ == '__main__':
    main()

